import torch.nn

import shared.torchtools


class PinnOfNavierStokesEquation(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.network: torch.nn.Sequential = torch.nn.Sequential(
            torch.nn.Linear(3, 300),
            torch.nn.Tanh(),
            *[
                 torch.nn.Linear(300, 300),
                 torch.nn.Tanh()
             ] * 3,
            torch.nn.Linear(300, 2)
        )
        self.lambda1: torch.nn.Parameter = torch.nn.Parameter(
            torch.tensor(0.), requires_grad=True)
        self.lambda2: torch.nn.Parameter = torch.nn.Parameter(
            torch.tensor(0.), requires_grad=True)

    def forward(self, t: torch.Tensor, x: torch.Tensor, y: torch.Tensor):
        assert t.shape == x.shape == y.shape
        assert len(t.shape) == 1

        if not x.requires_grad:
            x = shared.torchtools.clone_detached(x, True)
        if not y.requires_grad:
            y = shared.torchtools.clone_detached(y, True)
        txy = torch.stack([t, x, y], dim=-1)

        f = self.network(txy)
        psi = f[:, 0]
        p = f[:, 1]

        v, u = shared.torchtools.grad(psi, [x, y], True, True)
        v = -v

        return torch.stack([u, v, p], dim=-1)

    def physics_loss(self, point_count: int = 1000) -> torch.Tensor:
        t = shared.torchtools.rand(
            point_count, 0, 20, True)
        x = shared.torchtools.rand(
            point_count, 1, 8, True)
        y = shared.torchtools.rand(
            point_count, -2, 2, True)

        uvp = self(t, x, y)
        u = uvp[:, 0]
        v = uvp[:, 1]
        p = uvp[:, 2]

        du_dt, du_dx, du_dy = shared.torchtools.grad(u, [t, x, y], True, True)
        dv_dt, dv_dx, dv_dy = shared.torchtools.grad(v, [t, x, y], True, True)
        dp_dx, dp_dy = shared.torchtools.grad(p, [x, y], True, True)
        du_dx_dx, = shared.torchtools.grad(du_dx, [x], True, True)
        du_dy_dy, = shared.torchtools.grad(du_dy, [y], True, True)
        dv_dx_dx, = shared.torchtools.grad(dv_dx, [x], True, True)
        dv_dy_dy, = shared.torchtools.grad(dv_dy, [y], True, True)

        f = (du_dt + self.lambda1 * (u * du_dx + v * du_dy) +
             dp_dx - self.lambda2 * (du_dx_dx + du_dy_dy))
        g = (dv_dt + self.lambda1 * (u * dv_dx + v * dv_dy) +
             dp_dy - self.lambda2 * (dv_dx_dx + dv_dy_dy))

        return torch.mean(f ** 2 + g ** 2)
